#coding=utf-8
import json
import os
import smtplib
import datetime
from PIL import Image
from email.header import Header
from email.mime.image import MIMEImage
from email.mime.multipart import MIMEMultipart
from email.mime.text import MIMEText
from email.utils import parseaddr, formataddr
from pyecharts.charts import Bar, Page
from pyecharts import options as opts
# 导入输出图片工具
from pyecharts.render import make_snapshot
# 使用snapshot-selenium 渲染图片
from snapshot_selenium import snapshot


#自定义处理邮件收发地址的显示内容
def _format_addr(s):
    name, addr = parseaddr(s)
    # 将邮件的name转换成utf-8格式，addr如果是unicode，则转换utf-8输出，否则直接输出addr
    # return formataddr((\
    #     Header(name, 'utf-8').encode(),\
    #     addr.encode("utf-8") if isinstance(addr, bytes) else addr))
    return formataddr((Header(name, 'utf-8').encode(), addr))


def update_config_data():
    # 读取上次json数据
    with open('config.json', 'r', encoding='utf-8') as f:
        json_dict = json.loads(f.read())
    # 格式化json数据，保证每组为最近6条数据
    json_dict['url_num_list'].append(param['url_num'])
    json_dict['url_num_list'].pop(0)
    json_dict['vul_num_list'].append(param['vul_num'])
    json_dict['vul_num_list'].pop(0)
    json_dict['warning'].append(param['warning_num'])
    json_dict['warning'].pop(0)
    json_dict['information_leakage'].append(param['information_leakage_num'])
    json_dict['information_leakage'].pop(0)
    json_dict['honeypot'].append(param['honeypot_num'])
    json_dict['honeypot'].pop(0)
    print('json_dict', json_dict)
    return json_dict


def update_chart_data():
    json_dict = update_config_data()
    chart_list = [{'title': '渗透测试与漏洞发现数量趋势图', 'y_axis1': ['渗透测试'], 'y_axis2': ['漏洞']},
                  {'title': '安全预警数量数量趋势图', 'y_axis1': ['0day漏洞'], 'y_axis2': []},
                  {'title': 'GitHub&Gitee信息泄露数量趋势图', 'y_axis1': ['信息泄露'], 'y_axis2': []},
                  {'title': '蜜罐安全报警数量趋势图', 'y_axis1': ['威胁访问'], 'y_axis2': []}]
    # 格式化生成图表所需的list数据chart_list
    chart_list[0]['y_axis1'].append(json_dict['url_num_list'])
    chart_list[0]['y_axis2'].append(json_dict['vul_num_list'])
    chart_list[1]['y_axis1'].append(json_dict['warning'])
    chart_list[2]['y_axis1'].append(json_dict['information_leakage'])
    chart_list[3]['y_axis1'].append(json_dict['honeypot'])
    print('chart_list', chart_list)
    return chart_list


def create_image():
    chart_list = update_chart_data()
    # 获取本周所在周的前6周的周数，组成list数据
    current_week = int(today.strftime('%W'))
    week_list = []
    for i in range(current_week-6, current_week):
        week_list.append('第{}周'.format(i))
    # 依据chart_list和week_list生成html格式的图表，然后利用make_snapshot导出png格式图表
    image_list = []
    for i in range(len(chart_list)):
        bar = Bar()
        #设置x轴
        bar.add_xaxis(xaxis_data=week_list)
        #设置y轴
        if chart_list[i].get('y_axis2') == []:
            bar.add_yaxis(series_name=chart_list[i].get('y_axis1')[0], y_axis=chart_list[i].get('y_axis1')[1])
        else:
            bar.add_yaxis(series_name=chart_list[i].get('y_axis1')[0], y_axis=chart_list[i].get('y_axis1')[1])
            bar.add_yaxis(series_name=chart_list[i].get('y_axis2')[0], y_axis=chart_list[i].get('y_axis2')[1])
        bar.set_global_opts(title_opts=opts.TitleOpts(title=chart_list[i].get('title')))
        #生成html文件
        image_name = "{}.png".format(chart_list[i].get('title'))
        make_snapshot(snapshot, bar.render(), image_name)
        image_list.append(image_name)

    # 由于默认生成的图片太大，整体缩小50%，保存为新的图片
    image_dict = {}
    a = 0
    for i in image_list:
        im = Image.open(i)
        (x, y) = im.size  # 读取图片大小
        new_x = 900
        new_y = 500
        out = im.resize((new_x, new_y), Image.ANTIALIAS)
        new_image_name = 'new_' + i
        out.save(new_image_name)
        image_dict[new_image_name] = 'image' + str(a)
        a += 1
    return image_dict


def send_email():
    image_dict = create_image()
    # 设置邮件头，邮件主题，邮件内容，并将图片插入内容中
    last_week_start = today - datetime.timedelta(days=today.weekday() + 7)
    last_week_end = today - datetime.timedelta(days=today.weekday() + 1)
    msg = MIMEMultipart()
    msg_str = """
    <h3>一、安全检测</h3>
    <h4>1. 安全测试</h4>
    <p style="text-indent:35px;">本周渗透测试{}个域名，共发现{}个漏洞。</p>
    <img src="cid:image0">
    <h4>2. 安全扫描</h4>
    <p style="text-indent:35px;">对阿里云，idc公网扫描，无异常。</p>
    <h3>二、通用型漏洞监控</h3>
    <p style="text-indent:35px;">本周通用型漏洞监控发现{}个漏洞</p>
    <p>{}</p>
    <img src="cid:image1">
    <h3>三、GitHub&Gitee信息泄露监控</h3>
    <p style="text-indent:35px;">本周{}个代码泄露问题。</p>
    <img src="cid:image2">
    <h3>四、安全事件应急响应</h3>
    <p style="text-indent:35px;">本周蜜罐{}个告警。</p>
    <img src="cid:image3">
    <h3>五、黑产信息监控</h3>
    <p style="text-indent:35px;">本周未发现与xx相关安全事件。</p>
    """.format(param['url_num'], param['vul_num'], param['warning_num'], param['vul_name'], param['information_leakage_num'], param['honeypot_num'])
    msg.attach(MIMEText(msg_str, 'html', 'utf-8'))
    # 插入当前文件夹中图片
    for k, v in image_dict.items():
        with open(k, 'rb') as f:
            pic = MIMEImage(f.read())
            pic.add_header('Content-ID', '<{}>'.format(v))
            msg.attach(pic)
    # 设置邮件头信息
    msg['From'] = _format_addr(u'技术安全 <%s>' % sender)
    msg['To'] = ','.join(receiver)
    msg['Cc'] = ','.join(cc)
    subject = '安全检测周报（{}-{}）'.format(last_week_start.strftime('%Y年%m月%d日'), last_week_end.strftime('%Y年%m月%d日'))
    msg['Subject'] = Header(subject, 'utf-8')
    print('subject', subject)
    # 根据邮箱配置，发送邮件
    try:
        # smtpObj = smtplib.SMTP(mail_host, 25)
        smtpObj = smtplib.SMTP_SSL(mail_host, 465)
        # smtpObj.set_debuglevel(1)
        smtpObj.login(mail_user, mail_pass)
        smtpObj.sendmail(sender, all_receiver, msg.as_string())
        smtpObj.quit()
        return 'success'
    except smtplib.SMTPException as err:
        print('error:', err)
        return err


def update_config():
    # 更新config.json数据
    json_dict = update_config_data()
    new_json_str = json.dumps(json_dict)
    with open('config.json', 'w') as f:
        f.write(new_json_str)


def delete_file():
    # 删除该文件夹下所有.png和.html文件
    filename_list = os.listdir('.')
    for i in filename_list:
        if i.endswith(".png") or i.endswith(".html"):
            os.remove(i)


if __name__ == '__main__':
    # mail_host = 'smtp.163.com'
    # mail_user = 'unique@163.com'
    # mail_pass = 'ISTQGPQPKZ'
    # sender = 'unique@163.com'
    # receiver = ['5442502@qq.com', 'unique@163.com']
    # cc = []

    mail_host = 'smtp.exmail.qq.com'
    mail_user = 'tech-security@xxx.com'
    mail_pass = 'sFv2U4yeN'
    sender = 'tech-security@xxx.com'
    receiver = ['shield.agent.list@xxx.com']
    cc = ['jeffy@xxx.com', 'huhf@xxx.com', 'yuanxl@xxx.com']

    all_receiver = receiver + cc
    today = datetime.datetime.today()
    # "vul_name": "1. runc 符号链接挂载与容器逃逸漏洞（CVE-2021-30465）</br>2. 用友NC BeanShell远程代码执行漏洞（CNVD-2021-30167）"
    param = {
        "url_num": 5,
        "vul_num": 20,
        "warning_num": 6,
        "vul_name": "1. OpenEXR多个高危漏洞通告</br>2. Tapestry 序列化漏洞通告（CVE-2021-27850）</br>3. Istio 敏感信息窃取漏洞通告（CVE-2021-34824）</br>4. 西部数据NAS设备网络攻击通告</br>5. Windows Print Spooler远程代码执行漏洞通告（CVE-2021-1675）</br>6. PowerISO 内存越界写漏洞通告（CVE-2021-21871）",
        "information_leakage_num": 0,
        "honeypot_num": 0
    }
    success_send = send_email()  # 发送邮件
    delete_file()  # 删除当前文件夹下的.png和.html文件
    if success_send == 'success':
        update_config()  # 更新config.json数据
        print('Mail sending success')
    else:
        print('Mail sending failed')
